function Res = bciCheckClassifier(trainOnRuns,N)
% Res = bciCheckClassifier(trainOnRuns,N)
%
% PURPOSE:
% Check the performance on a trainingset using cross-validation
%
% INPUT:
% trainOnRuns:  Vector of training run numbers to include
%               If omitted or empty all runs in memory will be used
%               (TEST IMPLEMENTATION) Cell array of filenames to include
%
% OUTPUT:
%
% REQUIREMENTS: fieldtrip,spider,chronux

% XXXXXXXX CR wrote it
% 20110426 JR included option to load runs from file. Might chrash if
%           multiple runs are in one file
% 20110524 CC re-included doPlot and added output that displays top 50%
%           frequencies and channels
% 20110530 CR moved doPlot and output that displays top 50%
%           frequencies and channels to bciPlotSVMweights
% 20111007 CR generalized class label handling (no more restricted to
%             movement conditions)

if exist('svm','file')~=2,
    error('Check toolbox paths');
end
%% INPUT checking and data preparation
% get the data
global GLOBALbci;
global GLOBALtrainDat;

if nargin < 2,
    N=[]; % empty N is default, which means leave one run out or 3-fold if 1 run is given
end

if nargin==0 || isempty(trainOnRuns) || isnumeric(trainOnRuns) % USE ALL RUNS AVAILABLE IN MEMORY OR THOSE STATED IN TRAINONRUNS
    if isempty(GLOBALbci) || isempty(GLOBALtrainDat);
        error('GLOBALbci AND/OR GLOBALtrainDat not loaded in memory. Quitting.')
    end
    if nargin==0 || isempty(trainOnRuns) %nothing was passed
        trainOnRuns=find(~isempty(GLOBALtrainDat{:}));
    end
elseif iscell(trainOnRuns) %filenames were passed
    for k=1:length(trainOnRuns)
        load(trainOnRuns{k})
        GLOBALbci=bci;
        GLOBALtrainDat{k}=Dat; %THIS ASSUMES THAT THERE IS ONLY ON TRAINING BLOCK IN THR STRUCTUTRE
        load([trainOnRuns{k} '_Epoch']);
        GLOBALtrainDat{k}.trainEpoch=trainEpoch;
    end
    trainOnRuns=1:length(trainOnRuns)
end

if length(trainOnRuns)==1 && isempty(N),
    N=3;
end

%% Classification 
% Cross-validation
if ~isempty(N), % N-fold cross validation with only one tranining run
    splittedDat=splitTrainDat(GLOBALtrainDat(trainOnRuns),N);
    label=[];
    prediction = [];
    for k=1:N,
        [curbci,trainDat]=bciGetTrainDat(GLOBALbci,splittedDat{setdiff(1:N,k)});
        curbci=bciTrainClassifier(curbci,trainDat);
        validLabel=find(sum(splittedDat{k}.badChanMap(curbci.goodChan,:),1)==0);
        pred = zeros(size(validLabel));
        label=[label splittedDat{k}.trainLabel(validLabel)];
        for t=1:length(validLabel),
            dat = bciTransformData(splittedDat{k}.trainEpoch(:,:,validLabel(t)),curbci,1);
            pred(t)=GLOBALbci.param.testFunc(curbci,dat(:,curbci.featureMask));
        end
        prediction=[prediction,pred];
    end
else
    %leave-one-run-out
    label=[];
    prediction = [];
    for k=1:length(trainOnRuns),
        [curbci,trainDat]=bciGetTrainDat(GLOBALbci,GLOBALtrainDat{setdiff(trainOnRuns,trainOnRuns(k))});
        curbci=bciTrainClassifier(curbci,trainDat);
        validLabel=find(sum(GLOBALtrainDat{trainOnRuns(k)}.badChanMap(curbci.goodChan,:),1)==0);
        pred = zeros(size(validLabel));
        label=[label GLOBALtrainDat{trainOnRuns(k)}.trainLabel(validLabel)];
        for t=1:length(validLabel),
            dat = bciTransformData(GLOBALtrainDat{trainOnRuns(k)}.trainEpoch(:,:,validLabel(t)),curbci,1);
            pred(t)=GLOBALbci.param.testFunc(curbci,dat(:,curbci.featureMask));
        end
        prediction=[prediction,pred];
    end
end

predAcc=sum(prediction==label)/length(label);
if nargout==0,
    fprintf('PredAcc: %3.2f\n',predAcc*100);
else
    Res.predAcc = predAcc;
    Res.conditionID=[];
    Res.recall=[];
    Res.precision=[];
end
for cond = GLOBALbci.eventsToClassify,    
    if any(label==cond),
        recall = sum(prediction==cond&label==cond)/sum(label==cond);
        precision = sum(prediction==cond&label==cond)/...
            (sum(prediction==cond&label==cond)+sum(prediction==cond&label~=cond));
        if nargout==0;
            fprintf('Recall Cond %i: %3.2f Precision Cond %i : %3.2f\n',...
                cond,recall*100,cond, precision*100);
        else
           Res.conditionID=[Res.conditionID cond];
           Res.recall=[Res.recall recall];
           Res.precision=[Res.precision precision];
        end
    end
end


function Dat = splitTrainDat(trainDat,N)

% determine number of epochs of each run
nTrials = zeros(1,length(trainDat));
for k=1:length(trainDat),
    nTrials(k) = trainDat{k}.nEpochs;    
end

% determine assignment of original runs to splitted dat cells
selector = ceil((1:sum(nTrials))/(sum(nTrials)/N));
sourceRun = zeros(size(selector));
count = 0;
for k=1:length(nTrials),
    sourceRun(count+1:count+nTrials(k)) = k;
    count = count + nTrials(k);
end

Dat=cell(1,N);

% collect the data
for run=1:N,
    % required to merge runs ??
    mRuns = unique(sourceRun(selector==run));
    tmpEpoch = [];
    tmpLabel=[];
    tmpBadChanMap = [];
    for k=1:length(mRuns),
        idx = selector(sourceRun==mRuns(k))==run;
        tmpEpoch = cat(3,tmpEpoch,trainDat{mRuns(k)}.trainEpoch(:,:,idx));
        tmpLabel = cat(2,tmpLabel,trainDat{mRuns(k)}.trainLabel(idx));
        tmpBadChanMap = cat(2,tmpBadChanMap,trainDat{mRuns(k)}.badChanMap(:,idx));
    end
    Dat{run}.trainEpoch=tmpEpoch;
    Dat{run}.trainLabel=tmpLabel;
    Dat{run}.badChanMap=tmpBadChanMap;
end
